"""将用于风格迁移的两个文件夹内文件数量保持一致"""

import os
import random
import shutil
import cv2
from skimage import exposure

dataset_dir = "/disk527/sdb1/a804_cbf/datasets/crater_cyclegan/testA"
transfer_dir = "/disk527/sdb1/a804_cbf/datasets/chang_e/nac"
output_dir = "/disk527/sdb1/a804_cbf/datasets/crater_cyclegan/testB"
file_num = 14994
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

files = os.listdir(dataset_dir)
random.shuffle(files)
for it, file in enumerate(files[:file_num]):
    # shutil.copyfile(
    #     os.path.join(dataset_dir, file),
    #     os.path.join(output_dir, f"{it//102*50}_{it%102*50}.png"),
    # )
    img = cv2.imread(os.path.join(dataset_dir, file), cv2.IMREAD_GRAYSCALE)
    transfer = cv2.imread(os.path.join(transfer_dir, file), cv2.IMREAD_GRAYSCALE)
    multi = True if img.shape[-1] > 1 else False
    matched = exposure.match_histograms(transfer, img)
    # cv2.imwrite(os.path.join(output_dir, f"{it//114*50}_{it%114*50}.png"), img)
    cv2.imwrite(os.path.join(output_dir, file), matched)
